# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Staging Tranformation PIMD integrator
"""

from typing import Tuple

import torch
import math
from torch import Tensor

from .integrator import Integrator, _integrator_register
from ...system import Molecule
from ...function import get_arguments, get_tensor


@_integrator_register('stgpimd')
class PIMDMM(Integrator):
    r"""A leap-frog integrator based on "middle scheme" developed by Jian Liu, et al. It is a subclass of `Integrator`.

    Reference:

        Zhang, Z.; Yan, K; Liu, X.; Liu, J.
        A Leap-Frog Algorithm-based Efficient Unified Thermostat Scheme for Molecular Dynamics [J].
        Chinese Science Bulletin, 2018, 63(33): 3467-3483.

    Args:

        system (Molecule):          Simulation system

        thermostat (Thermostat):    Thermostat for temperature coupling. Default: None

        barostat (Barostat):        Barostat for pressure coupling. Default: None

        constraint (Constraint):    Constraint algorithm. Default: None


    Supported Platforms:

        ``Ascend`` ``GPU``

    """
    def __init__(self,
                 system: Molecule,
                 qm_index,
                 temperature = 300.,
                 gamma0 = 10.,
                 remove_com = True,
                 **kwargs
                 ):

        super().__init__(
            system=system,
            thermostat=None,
            barostat=None,
            constraint=None,
        )
        self._kwargs = get_arguments(locals(), kwargs)
        self.name = 'PILE'
        n = self.num_walker
        self.temperature = self._get_mw_tensor(temperature, 'temperature')
        self.qm_index = get_tensor(qm_index).reshape(-1)
        self._total_index = torch.arange(system.coordinate.shape[1])
        self.mm_index = self._total_index[~torch.isin(self._total_index,self.qm_index)]


        boltzmann = system.units.boltzmann
        hbar = 1.054572e-37 * 6.02e23 * 1e12
        self.beta = 1. / (boltzmann * self.temperature)
        self.omega_n = math.sqrt(n) / (self.beta * hbar) *torch.ones(n,1,1)

        self.gamma = math.sqrt(n) / (self.beta * hbar) *torch.ones(n,1,1)
        self.gamma0 = gamma0
        self.gamma[0] = self.gamma0
        self.c1 = torch.exp(-self.gamma*self.time_step)
        self.c2 = torch.sqrt(1-self.c1**2) / self.kinetic_unit_scale

        self.m_bar = torch.zeros_like(self.atom_mass[:,self.qm_index].unsqueeze(-1)).repeat(n,1,1)
        self.m_ratio = torch.cat([torch.ones(1,), torch.arange(2,n+1) / torch.arange(1,n)]).unsqueeze(-1).unsqueeze(-1)
        self.m_bar[1:] = self.m_ratio[1:] * self.atom_mass[:,self.qm_index].unsqueeze(-1)
        self.m_tilde = self.m_ratio * self.atom_mass[:,self.qm_index].unsqueeze(-1)
        self._mm_inv_mass = system.inv_mass[:,self.mm_index].unsqueeze(-1)
        self._mm_inv_sqrt_mass = torch.sqrt(self._mm_inv_mass)
        self._mm_mass = self._atom_mass[:,self.mm_index]
        self._mm_system_mass = self._mm_mass.sum(dim=-2,keepdim=True)

        # \f = 1 - exp(-\gamma * dt)
        self.friction = 1.0 - torch.exp(-self.gamma0*self.time_step)
        # k = \sqrt(f * (2 - f) * k_B * T)
        self.random_const = self.friction * (2 - self.friction) * self.boltzmann / self.kinetic_unit_scale

        self._j = torch.arange(1,n+1).reshape(-1,1,1)
        self._ratio = (self._j-1) / self._j
        self.sqrt_m_tilde = torch.sqrt(self.m_tilde)
        self.remove_com = remove_com

        self._Fmatrix = torch.zeros((n,n))
        for i in range(n):
            for j in range(i + 1):
                if i == j:
                    self._Fmatrix[i, j] = 1
                else:
                    self._Fmatrix[i, j] = j / i
        self._Fmatrix[0] = torch.ones(n)
        self._Fmatrix /= n

    def print_info(self):
        info = {'Simulation Ensemble': 'NVT Ensemble',
                'Integrator Type': 'STG-PIMD / Leap-Frog',
                'time step': self.time_step,
                'temperature': self.temperature.item(),
                'gamma0': self.gamma0,}
        return info

    def forward(self,
               coordinate: Tensor,
               velocity: Tensor,
               force: Tensor,
               energy: Tensor,
               virial: Tensor = None,
               pbc_box: Tensor = None,
               steps: int = 0,
               **kwargs
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:

        qm_variables = self.PI_integrator(coordinate[:,self.qm_index],
                            velocity[:,self.qm_index],
                            force[:,self.qm_index],
                            energy,
                            virial,
                            pbc_box,
                            steps,
                            **kwargs)
        
        mm_variables = self.MM_integrator(coordinate[:,self.mm_index].mean(0, keepdim=True),
                            velocity[:,self.mm_index].mean(0, keepdim=True),
                            force[:,self.mm_index].mean(0, keepdim=True),
                            energy,
                            virial,
                            pbc_box,
                            steps,
                            **kwargs)
        
        coordinate_new = torch.cat([qm_variables['coordinate'],mm_variables['coordinate'].expand(coordinate.shape[0],-1,-1)],dim=1)
        velocity_new = torch.cat([qm_variables['velocity'],mm_variables['velocity'].expand(velocity.shape[0],-1,-1)],dim=1)
        force_new = torch.cat([qm_variables['force'],mm_variables['force'].expand(force.shape[0],-1,-1)],dim=1)
        
        variables = {'coordinate': coordinate_new,
                   'velocity': velocity_new,
                    'force': force_new,
                    'energy': qm_variables['energy']+mm_variables['energy'],
                    'virial': qm_variables['virial']+mm_variables['virial'],
                    'pbc_box': mm_variables['pbc_box'],
                    }
        
        return variables
 
    def PI_langevin(self, p, steps):
        if self.control_step == 1 or steps % self.control_step == 0:
            p = self.c1 * p + self.c2 * torch.sqrt(1 / self.beta) * self.sqrt_m_tilde * torch.randn(p.shape, device=p.device)
        return p
    
    def get_PI_com_velocity(self, momentum: Tensor):
        com_vel = momentum.sum(-2,keepdim=True)/self.m_tilde.sum(-2,keepdim=True) * self.m_tilde
        return com_vel
    
    def get_MM_com_velocity(self, velocity: Tensor, keepdims: bool = True) -> Tensor:

        # (B, A, D) * (B, A, 1)
        com_vel = velocity * self._mm_mass
        # (B, D) or (B, 1, D) <- (B, A, D)
        com_vel = torch.sum(com_vel, dim=-2, keepdim=keepdims)
        # (B, D) or (B, 1, D) = (B, D) or (B, 1, D) / (B, 1)
        com_vel = com_vel / self._mm_system_mass.sum(dim=-2, keepdim=keepdims)
        return com_vel
    
    def MM_langevin(self, velocity, steps):
        if self.control_step == 1 or steps % self.control_step == 0:
            velocity += -self.friction * velocity + torch.sqrt(self.random_const * self.temperature) * \
                    self._mm_inv_sqrt_mass * torch.randn_like(velocity)
        return velocity
    
    def MM_integrator(self,
               coordinate: Tensor,
               velocity: Tensor,
               force: Tensor,
               energy: Tensor,
               virial: Tensor = None,
               pbc_box: Tensor = None,
               steps: int = 0,
               **kwargs
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
        
         # (B,A,D) = (B,A,D) * (B,A,1)
        acceleration = self.acc_unit_scale * force * self._mm_inv_mass

        # v(t+0.5) = v(t-0.5) + a(t) * dt
        velocity_half = velocity + acceleration * self.time_step
        # (B,A,D) = (B,A,D) - (B,1,D)
        velocity_half = velocity_half - self.get_MM_com_velocity(velocity_half)

        # R(t+0.5) = R(t) + v(t+0.5) * dt
        coordinate_half = coordinate + velocity_half * self.time_step * 0.5

        variables = {'coordinate': coordinate_half,
                    'velocity': velocity_half,
                    'force': force,
                    'energy': energy,
                    'virial': virial,
                    'pbc_box': pbc_box,
                    }

        if self.temperature is not None:
            # v'(t+0.5) = f_T[v(t+0.5)]
            variables['velocity'] = self.MM_langevin(velocity_half, steps=steps)

        # R(t+1) = R(t+0.5) + v'(t+0.5) * dt
        variables['coordinate'] = variables['coordinate'] + variables['velocity'] * self.time_step * 0.5

        return variables
        
    def PI_integrator(self,
               coordinate: Tensor,
               velocity: Tensor,
               force: Tensor,
               energy: Tensor,
               virial: Tensor = None,
               pbc_box: Tensor = None,
               steps: int = 0,
               **kwargs
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
        # (B,A,D) = (B,A,D) * (B,A,1)
        x_tilde = coordinate - self._ratio * coordinate.roll(-1,dims=0) - coordinate[[0]] /  self._j
        x_tilde[0] = coordinate[0]
        f_tilde = torch.einsum('jk,kad->jad',self._Fmatrix,force)

        p_tilde = velocity * self.m_tilde / torch.sqrt(self.m_ratio)
        p_tilde = p_tilde + self.acc_unit_scale * f_tilde * self.time_step - self.omega_n ** 2 * x_tilde * self.time_step * self.m_bar
        # (B,A,D) = (B,A,D) - (B,1,D)
        if self.remove_com:
            p_tilde = p_tilde - self.get_PI_com_velocity(p_tilde)

        # R(t+0.5) = R(t) + v(t+0.5) * dt
        x_tilde = x_tilde + p_tilde / self.m_tilde * 0.5 * self.time_step

        if self.temperature is not None:
            # v'(t+0.5) = f_T[v(t+0.5)]
            p_tilde = self.PI_langevin(p_tilde, steps)

        # R(t+1) = R(t+0.5) + v'(t+0.5) * dt
        x_tilde = x_tilde + p_tilde / self.m_tilde * 0.5 * self.time_step
        velocity = p_tilde / self.m_tilde * torch.sqrt(self.m_ratio)

        summation_x = x_tilde / (self._j-1)
        summation_x[0] = 0
        coordinate = x_tilde[[0]] + (self._j-1) * torch.flip(torch.cumsum(torch.flip(summation_x,dims=[0]),0),dims=[0])

        variables = {'coordinate': coordinate,
                    'velocity': velocity,
                    'force': force,
                    'energy': energy,
                    'virial': virial,
                    'pbc_box': pbc_box,
                    }
        return variables


        



